import os
import pydicom
import shutil

def classify_series_by_window(input_dir, output_dir):
    """
    将每个序列文件夹分类保存：
    - WindowWidth 在 [1000,2000] 且 WindowCenter 在 [-600,-400] → Class1
    - 其他情况 → Class2
    """
    class1_dir = os.path.join(output_dir, "fc")
    class2_dir = os.path.join(output_dir, "zg")
    os.makedirs(class1_dir, exist_ok=True)
    os.makedirs(class2_dir, exist_ok=True)

    # 遍历序列文件夹（每个文件夹对应一个序列）
    for series_folder in os.listdir(input_dir):
        series_path = os.path.join(input_dir, series_folder)
        if not os.path.isdir(series_path):
            continue

        # 随便取序列里的一个dicom文件来读取 Window 信息
        dicom_files = [f for f in os.listdir(series_path) if f.lower().endswith(".dcm")]
        if not dicom_files:
            continue

        dicom_path = os.path.join(series_path, dicom_files[0])

        try:
            ds = pydicom.dcmread(dicom_path, stop_before_pixels=True)

            wc = ds.get("WindowCenter", None)
            ww = ds.get("WindowWidth", None)

            # WindowCenter/Width 可能是数组，取第一个
            if isinstance(wc, (list, pydicom.multival.MultiValue)):
                wc = float(wc[0])
            elif wc is not None:
                wc = float(wc)

            if isinstance(ww, (list, pydicom.multival.MultiValue)):
                ww = float(ww[0])
            elif ww is not None:
                ww = float(ww)

            # 判断条件
            if wc is not None and ww is not None and (1000 <= ww <= 2000) and (-600 <= wc <= -400):
                target_dir = os.path.join(class1_dir, series_folder)
            else:
                target_dir = os.path.join(class2_dir, series_folder)

            # 拷贝整个序列目录
            if not os.path.exists(target_dir):
                shutil.copytree(series_path, target_dir)

        except Exception as e:
            print(f"跳过 {series_path}, 错误: {e}")


if __name__ == "__main__":
    input_dir = r"/media/imed/d9f520b2-6e7a-4842-af77-fba3ad2d8d27/heart_mao_200_123/123CT_split"   # 已经按序列分好的文件夹
    output_dir = r"/media/imed/d9f520b2-6e7a-4842-af77-fba3ad2d8d27/heart_mao_200_123/123CT_split_classified"        # 分类后的输出目录
    classify_series_by_window(input_dir, output_dir)